package com.octopus.core.limiter;

import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

/**
 * 速率限制器，主要是为了限制页面爬虫爬取速率限制，防止爬虫被目标服务器限制。
 *
 * <p>每隔一段时间会释放许可信号，爬虫引擎获取许可后方可执行采取任务。
 *
 * @author shoulai.yang@gmail.com
 * @date 2019/07/31
 */
public final class RateLimiter {

  /** 一段时间内最大执行任务数 */
  private int maxPermits;

  /** 信号量 */
  private Semaphore semaphore;

  /** 许可信号量释放执行线程池 */
  private ScheduledExecutorService scheduler;

  /** 时间单元 */
  private TimeUnit timePeriod;

  public RateLimiter(int permits, TimeUnit unit) {
    this.scheduler =
        new ScheduledThreadPoolExecutor(
            1,
            r -> {
              Thread thread = new Thread(r);
              thread.setName("RateLimiter");
              return thread;
            });
    this.maxPermits = permits;
    this.semaphore = new Semaphore(this.maxPermits);
    this.timePeriod = unit;
  }

  /**
   * 阻塞获取许可
   *
   * @throws InterruptedException 线程中断异常
   */
  public void acquire() throws InterruptedException {
    this.semaphore.acquire();
  }

  /** 开启定期释放许可信号的线程 */
  void start() {
    this.scheduler.scheduleAtFixedRate(
        () -> {
          this.semaphore.release(this.maxPermits - this.semaphore.availablePermits());
        },
        1,
        1,
        this.timePeriod);
  }

  /** 停止释放许可信号 */
  void stop() {
    this.scheduler.shutdownNow();
  }
}
